#include "stdafx.h"
#include "FourierTransformFactoryFFTW.h"
#include "includeIpp.h"
#include "IppUtils.h"

#include "include_fftw.h"
#include "fftutils.h"
#include "delayloadUtils.h"

using namespace FTLib;

// TODO: no, the packed layout of IPP is different from FFTW, we need an extra interface for this...
class CRealToComplex_PackedFftwComplexFFTWFft : public IRealToComplex_PackedFftwComplex
{
protected:
	uint32_t width, height;
	fftwf_plan	fftwPlan;
	float normalizationFactor;
public:
	CRealToComplex_PackedFftwComplexFFTWFft() = delete;
	CRealToComplex_PackedFftwComplexFFTWFft(uint32_t width, uint32_t height, double normalization) : width(width), height(height)
	{
		if (IsNormalizationFactorDivBySize(normalization))
		{
			this->normalizationFactor = (float)(this->width*this->height);
		}
		else if (IsNormalizationFactorDivBySqrtSize(normalization))
		{
			this->normalizationFactor = sqrtf((float)(this->width*this->height));
		}
		else
		{
			this->normalizationFactor = (float)normalization;
		}

		uint32_t destHeight = height;
		uint32_t destWidth = width / 2 + 1;
		size_t sizeDestBuffer = destWidth*destHeight*sizeof(float) * 2;
		float* ptrTempDestBuffer = (float*)_aligned_malloc(sizeDestBuffer, 32);

		int n[2] = { width, height };

		// It seems... that if we use FFTW_ESTIMATE, then we have to provide a destination-buffer but no source buffer
		// (no crash, but wrong result if destination-buffer is also null). On the other hand, with FFTW_ESTIMATE we
		// need both source and destination. 
		// Also, it seems that after plan-generation we can delete the buffers.
		fftwf_plan plan = fftwf_plan_dft_r2c(2, n, nullptr, (fftwf_complex*)ptrTempDestBuffer, /*FFTW_MEASURE*/FFTW_ESTIMATE);
		this->fftwPlan = plan;
		ZeroMemory(ptrTempDestBuffer, sizeDestBuffer);
		_aligned_free(ptrTempDestBuffer);
	}

	~CRealToComplex_PackedFftwComplexFFTWFft()
	{
		fftwf_destroy_plan(this->fftwPlan);
	}

	virtual void RealToComplex_PackedFftwComplex(const FourierTransfRealInputData2d& input, const FourierTransfComplexOutput2d& output, const FourierTransformAdditional* pAdditional)
	{
		// TODO: Deal with different strides...

		//int n[2] = { input.width, input.height };
		//fftwf_plan plan = fftwf_plan_dft_r2c(2, n, const_cast<float*>(input.ptrData), (fftwf_complex*)output.ptrData, FFTW_ESTIMATE);

		//fftwf_execute(plan);// , const_cast<float*>(input.ptrData), (fftwf_complex*)output.ptrData);
		fftwf_execute_dft_r2c(this->fftwPlan, const_cast<float*>(input.ptrData), (fftwf_complex*)output.ptrData);

		if (this->normalizationFactor != 1)
		{
			IppiSize ippiSize = { (int)(output.width /*/ 2 + 1*/) * 2, (int)output.height };
			IppStatus status = ippiDivC_32f_C1IR(this->normalizationFactor, (float*)output.ptrData, /*(output.width / 2 + 1) * sizeof(float) * 2*/output.pitch, ippiSize);
		}

		//int n[2] = { width, height };
		//fftwf_plan plan = fftwf_plan_dft_r2c(2, n, const_cast<float*>(input.ptrData), (fftwf_complex*)output.ptrData, FFTW_ESTIMATE);
		//fftwf_execute_dft_r2c(plan, const_cast<float*>(input.ptrData), (fftwf_complex*)output.ptrData);
	}
};

class CRealToComplex_ComplexFFTWFft : public IRealToComplex_Complex, protected CRealToComplex_PackedFftwComplexFFTWFft
{
public:
	CRealToComplex_ComplexFFTWFft(uint32_t width, uint32_t height, double normalization) : CRealToComplex_PackedFftwComplexFFTWFft(width, height, normalization)
	{}

	virtual void RealToComplex_Complex(const FourierTransfRealInputData2d& input, const FourierTransfComplexOutput2d& output, const FourierTransformAdditional* pAdditional)
	{
		FourierTransfComplexOutput2d outputPacked;
		outputPacked.width = this->width / 2 + 1;
		outputPacked.height = this->height;
		size_t sizeDestBuffer = outputPacked.width*outputPacked.height *sizeof(float) * 2;
		outputPacked.ptrData = (float*)_aligned_malloc(sizeDestBuffer, 32);
		outputPacked.pitch = outputPacked.width*sizeof(float) * 2;

		this->RealToComplex_PackedFftwComplex(input, outputPacked, pAdditional);
		ConvertFftwToComplex(input.width, input.height, outputPacked.ptrData, output.ptrData, output.pitch);
		_aligned_free(outputPacked.ptrData);
	}
};

class CFftwUtils
{
public:
	void ThrowIfStrideIsNotAcceptable(uint32_t stride, uint32_t elementSize, uint32_t width)
	{
		if ((stride % elementSize) != 0)
		{
			throw std::invalid_argument("Invalid stride");
		}

		if (stride < width*elementSize)
		{
			throw std::invalid_argument("Invalid stride");
		}
	}
};

template<int tSign>
class CComplexToComplexBase_ComplexFFTWFft : protected CFftwUtils
{
private:
	uint32_t width, height;
	uint32_t planStrideSrc, planStrideDest;
	fftwf_plan	fftwPlan;
	bool planCreated;
	float normalizationFactor;
protected:
	CComplexToComplexBase_ComplexFFTWFft(uint32_t width, uint32_t height, double normalization) : CComplexToComplexBase_ComplexFFTWFft(width, height, normalization, false)
	{}

	CComplexToComplexBase_ComplexFFTWFft(uint32_t width, uint32_t height, double normalization, bool preCreate) : CComplexToComplexBase_ComplexFFTWFft(width, height, normalization, CalcDefaultStride(width), preCreate)
	{}

	CComplexToComplexBase_ComplexFFTWFft(uint32_t width, uint32_t height, double normalization, uint32_t stride, bool preCreate) : CComplexToComplexBase_ComplexFFTWFft(width, height, normalization, stride, stride, preCreate)
	{}

	CComplexToComplexBase_ComplexFFTWFft(uint32_t width, uint32_t height, double normalization, uint32_t strideSrc, uint32_t strideDst, bool preCreate) : planCreated(false), width(width), height(height), planStrideSrc(strideSrc), planStrideDest(strideDst)
	{
		if (preCreate == true)
		{
			this->CreatePlan(nullptr, nullptr);
			this->planCreated = true;
		}

		if (IsNormalizationFactorDivBySize(normalization))
		{
			this->normalizationFactor = (float)(this->width*this->height);
		}
		else if (IsNormalizationFactorDivBySqrtSize(normalization))
		{
			this->normalizationFactor = sqrtf((float)(this->width*this->height));
		}
		else
		{
			this->normalizationFactor = (float)normalization;
		}
	}

	~CComplexToComplexBase_ComplexFFTWFft()
	{
		this->DeletePlan();
	}

	void CoreComplexToComplex_Complex(const FourierTransfComplexInput2d& input, const FourierTransfComplexOutput2d& output)
	{
		if (this->planCreated != true)
		{
			this->planStrideSrc = input.pitch;
			this->planStrideDest = output.pitch;
			this->CreatePlan(input.ptrData, output.ptrData);

			// it seems to work when passing in null - at least it seems...
			//this->CreatePlan(0, 0);	
		}

		if (input.pitch != this->planStrideSrc || output.pitch != this->planStrideDest)
		{
			// well, sorry, then we have to delete the plan and re-create it
			this->planStrideSrc = input.pitch;
			this->planStrideDest = output.pitch;
			this->CreatePlan(input.ptrData, output.ptrData);

		}
#if 0
		int n[2] = { this->width, this->height };
		//const int* inembed = nullptr;// n;
		//const int* onembed = nullptr;// n;
		//int strideIn = 2;
		//int strideOut = 2;

		// TODO: This seems to be the way how to pass in the stride (our concept of it).
		//       If the stride is not a multiple of 4 then we have to bail out (or make a copy of the data).
		int inembed[2] = { 2, input.pitch / (2 * sizeof(float)) };
		int onembed[2] = { 2, output.pitch / (2 * sizeof(float)) };

		//fftwf_plan plan = fftwf_plan_dft_2d(2, 2, (fftwf_complex*)input.ptrData, (fftwf_complex*)output.ptrData, FFTW_FORWARD, FFTW_ESTIMATE);
		fftwf_plan plan = fftwf_plan_many_dft(2, n, 1, (fftwf_complex*)input.ptrData, inembed/*0*/, 1, 1, (fftwf_complex*)output.ptrData, onembed/*0*/, 1, 1, FFTW_FORWARD, FFTW_ESTIMATE);

		//fftwf_plan plan = fftwf_plan_many_dft(2, n, 1, (fftwf_complex*)input.ptrData, inembed, strideIn, 0, (fftwf_complex*)output.ptrData, onembed, strideOut, 0, FFTW_FORWARD, FFTW_ESTIMATE);

		//fftwf_execute_dft(plan, (fftwf_complex*)input.ptrData, (fftwf_complex*)output.ptrData);
#endif
		fftwf_execute_dft(this->fftwPlan, (fftwf_complex*)input.ptrData, (fftwf_complex*)output.ptrData);

		this->ApplyNormalization(output);


		/*FourierTransfComplexOutput2d outputPacked;
		outputPacked.width = this->width / 2 + 1;
		outputPacked.height = this->height;
		size_t sizeDestBuffer = outputPacked.width*outputPacked.height *sizeof(float) * 2;
		outputPacked.ptrData = (float*)_aligned_malloc(sizeDestBuffer, 32);
		outputPacked.pitch = outputPacked.width*sizeof(float) * 2;

		this->RealToComplex_PackedFftwComplex(input, outputPacked);
		ConvertFftwToComplex(input.width, input.height, outputPacked.ptrData, output.ptrData, output.pitch);
		_aligned_free(outputPacked.ptrData);*/
	}

private:
	void ApplyNormalization(const FourierTransfComplexOutput2d& output)
	{
		if (this->normalizationFactor != 1)
		{
			IppiSize size = { 2 * output.width, output.height };
			ippiDivC_32f_C1IR(this->normalizationFactor, output.ptrData, output.pitch, size);
		}
	}

	bool CreatePlan(const void* ptrSrc, void* ptrDst)
	{
		this->ThrowIfStrideIsNotAcceptable(this->planStrideSrc, 2 * sizeof(float), this->width);
		this->ThrowIfStrideIsNotAcceptable(this->planStrideDest, 2 * sizeof(float), this->width);
		int n[2] = { this->width, this->height };

		// TODO: deal with the case "stride is not a multiple of four" (which probably means to make a temp copy)
		// This seems to be the way how to pass in the stride (our concept of it).
		int inembed[2] = { 2, this->planStrideSrc / (2 * sizeof(float)) };
		int onembed[2] = { 2, this->planStrideDest / (2 * sizeof(float)) };

		this->fftwPlan = fftwf_plan_many_dft(2, n, 1, (fftwf_complex*)ptrSrc/*input.ptrData*/, inembed/*0*/, 1, 1, (fftwf_complex*)ptrDst/*output.ptrData*/, onembed/*0*/, 1, 1, tSign/*FFTW_FORWARD*/, FFTW_ESTIMATE);
		this->planCreated = true;
		return true;
	}

	void DeletePlan()
	{
		fftwf_destroy_plan(this->fftwPlan);
		this->fftwPlan = nullptr;
		this->planCreated = false;
	}

	static uint32_t CalcDefaultStride(uint32_t width)
	{
		return width * 2 * sizeof(float);
	}
};

class CComplexToComplexFwd_ComplexFFTWFft : public CComplexToComplexBase_ComplexFFTWFft < FFTW_FORWARD >, public IComplexToComplexFwd_Complex
{
public:
	CComplexToComplexFwd_ComplexFFTWFft(uint32_t width, uint32_t height, double normalization) : CComplexToComplexBase_ComplexFFTWFft<FFTW_FORWARD>(width, height, normalization)
	{}

	CComplexToComplexFwd_ComplexFFTWFft(uint32_t width, uint32_t height, double normalization, bool preCreate) : CComplexToComplexBase_ComplexFFTWFft<FFTW_FORWARD>(width, height, normalization, preCreate)
	{}

	CComplexToComplexFwd_ComplexFFTWFft(uint32_t width, uint32_t height, double normalization, uint32_t stride, bool preCreate) : CComplexToComplexBase_ComplexFFTWFft<FFTW_FORWARD>(width, height, normalization, stride, stride, preCreate)
	{}

	CComplexToComplexFwd_ComplexFFTWFft(uint32_t width, uint32_t height, double normalization, uint32_t strideSrc, uint32_t strideDst, bool preCreate) : CComplexToComplexBase_ComplexFFTWFft<FFTW_FORWARD>(width, height, normalization, strideSrc, strideDst, preCreate)
	{}

	virtual void ComplexToComplex_Complex(const FourierTransfComplexInput2d& input, const FourierTransfComplexOutput2d& output, const FourierTransformAdditional* pAdditional)
	{
		this->CoreComplexToComplex_Complex(input, output);
	}
};

class CComplexToComplexInv_ComplexFFTWFft : public CComplexToComplexBase_ComplexFFTWFft < FFTW_BACKWARD>, public IComplexToComplexInv_Complex
{
public:
	CComplexToComplexInv_ComplexFFTWFft(uint32_t width, uint32_t height, double normalization) : CComplexToComplexBase_ComplexFFTWFft<FFTW_BACKWARD>(width, height, normalization)
	{}

	CComplexToComplexInv_ComplexFFTWFft(uint32_t width, uint32_t height, double normalization, bool preCreate) : CComplexToComplexBase_ComplexFFTWFft<FFTW_BACKWARD>(width, height, normalization, preCreate)
	{}

	CComplexToComplexInv_ComplexFFTWFft(uint32_t width, uint32_t height, double normalization, uint32_t stride, bool preCreate) : CComplexToComplexBase_ComplexFFTWFft<FFTW_BACKWARD>(width, height, normalization, stride, stride, preCreate)
	{}

	CComplexToComplexInv_ComplexFFTWFft(uint32_t width, uint32_t height, double normalization, uint32_t strideSrc, uint32_t strideDst, bool preCreate) : CComplexToComplexBase_ComplexFFTWFft<FFTW_BACKWARD>(width, height, normalization, strideSrc, strideDst, preCreate)
	{}

	virtual void ComplexToComplex_Complex(const FourierTransfComplexInput2d& input, const FourierTransfComplexOutput2d& output, const FourierTransformAdditional* pAdditional)
	{
		this->CoreComplexToComplex_Complex(input, output);
	}
};

class CComplexHermetianToReal_ComplexFFTWFft : public IComplexHermitianToReal
{
protected:
	uint32_t width, height;
	fftwf_plan	fftwPlan;
	float normalizationFactor;
public:
	CComplexHermetianToReal_ComplexFFTWFft() = delete;
	CComplexHermetianToReal_ComplexFFTWFft(uint32_t width, uint32_t height, double normalization) : width(width), height(height)
	{
		if (IsNormalizationFactorDivBySize(normalization))
		{
			this->normalizationFactor = (float)(this->width*this->height);
		}
		else if (IsNormalizationFactorDivBySqrtSize(normalization))
		{
			this->normalizationFactor = sqrtf((float)(this->width*this->height));
		}
		else
		{
			this->normalizationFactor = (float)normalization;
		}

		/*uint32_t destHeight = height;
		uint32_t destWidth = width / 2 + 1;*/
		size_t sizeDestBuffer = width*height*sizeof(float);// destWidth*destHeight*sizeof(float) * 2;
		float* ptrTempDestBuffer = (float*)_aligned_malloc(sizeDestBuffer, 32);

		int n[2] = { width, height };
		fftw_r2r_kind fftwKind[2] = { FFTW_HC2R, FFTW_HC2R };

		// It seems... that if we use FFTW_ESTIMATE, then we have to provide a destination-buffer but no source buffer
		// (no crash, but wrong result if destination-buffer is also null). On the other hand, with FFTW_ESTIMATE we
		// need both source and destination. 
		// Also, it seems that after plan-generation we can delete the buffers.
		//fftwf_plan plan = fftwf_plan_r2r(2, n, nullptr, (float*)ptrTempDestBuffer, fftwKind, /*FFTW_MEASURE*/FFTW_ESTIMATE);
		//fftwf_plan plan = fftwf_plan_dft_c2r_2d(width, height, nullptr, (float*)ptrTempDestBuffer, FFTW_ESTIMATE);

		//int n[2] = { this->width, this->height };

		// TODO: deal with the case "stride is not a multiple of four" (which probably means to make a temp copy)
		// This seems to be the way how to pass in the stride (our concept of it).

		int inembed[2] = { 1, width };	// this is what we need to set when the source is the "complete complex matrix" -> stride is "width"
		// which is a the number of "complex floats"
		int onembed[2] = { 1, width };

		//void* ptrTempSrc = _aligned_malloc(width*height*sizeof(float) * 2, 32);

		fftwf_plan plan = fftwf_plan_many_dft_c2r(2, n, 1, (fftwf_complex*)nullptr/*ptrTempSrc*//*input.ptrData*/,
			inembed/*0*/, 1, 1, (float*)ptrTempDestBuffer/*output.ptrData*/, onembed/*0*/, 1, 1, FFTW_ESTIMATE);


		this->fftwPlan = plan;
		//ZeroMemory(ptrTempDestBuffer, sizeDestBuffer);
		_aligned_free(ptrTempDestBuffer);
	}

	~CComplexHermetianToReal_ComplexFFTWFft()
	{
		fftwf_destroy_plan(this->fftwPlan);
	}

	virtual void ComplexHermitianToReal(const FourierTransfComplexInput2d& input, const FourierTransfRealOutputData2d& output, const FourierTransformAdditional* pAdditional)
	{
		/*static const float factors[] =
		{ 118.000000, 0.000000000, -44.0000000, 62.3538017,
		-50.0000000, -71.0140991, 76.0000000, 1.73204994,
		-50.0000000, 71.0140991, -36.5000000, -66.6839981 };*/
		//fftwf_execute_dft_c2r(this->fftwPlan, (fftwf_complex*)/*input.ptrData*/factors, (float*)output.ptrData);
		//memcpy((void*)input.ptrData, factors, sizeof(float) * 12);

		// TODO: for whatever reasons, the input is overwritten - no idea whether it is a feature or a bug
		//        and so far I couldn't figure out how to do it properly. So... either fix this or
		//        we need to copy the input (maybe we can do an inplace-operation?)

		fftwf_execute_dft_c2r(this->fftwPlan, (fftwf_complex*)input.ptrData, (float*)output.ptrData);

		if (this->normalizationFactor != 1)
		{
			IppiSize size = { output.width, output.height };
			ippiDivC_32f_C1IR(this->normalizationFactor, output.ptrData, output.pitch, size);
		}
	}

};

//------------------------------------------------------------------------------

/*static*/ const char* CFourierTransformFactoryFFTWFft::DllNameFFTWF = "libfftw3f-3.dll";

/*static*/ bool CFourierTransformFactoryFFTWFft::CanLoadLibrary()
{
	return TryToLoadAllImports(CFourierTransformFactoryFFTWFft::DllNameFFTWF);
}

/*static*/ bool CFourierTransformFactoryFFTWFft::CanOperate()
{
	return true;
}

/*static*/ void CFourierTransformFactoryFFTWFft::GetVersionString(std::wstring& str)
{
	if (!GetIsAvailable())
	{
		str = L"ERROR";
		return;
	}

	// we need to go for the version by ourselves, since they are not functions but exported statics
	// (for which delay-loading does not work)
	const char* pVersion;
	const char* pCc;
	pVersion = pCc = "error";
	HMODULE hMod = GetModuleHandleA(CFourierTransformFactoryFFTWFft::DllNameFFTWF);
	if (hMod != NULL)
	{
		FARPROC fpFFTWF_Version = GetProcAddress(hMod, "fftwf_version");
		if (fpFFTWF_Version != NULL)
		{
			pVersion = (const char*)fpFFTWF_Version;
		}

		FARPROC fpFFTWF_Cc = GetProcAddress(hMod, "fftwf_cc");
		if (fpFFTWF_Cc != NULL)
		{
			pCc = (const char*)fpFFTWF_Cc;
		}
	}

	std::stringstream ss;
	ss << R"(version=")" << pVersion << R"(", cc=")" << pCc << R"(")";

	std::wstring_convert<std::codecvt_utf8_utf16<wchar_t>> converter;
	str = converter.from_bytes(ss.str().c_str());
}

/*virtual*/std::shared_ptr<IRealToComplex_PackedFftwComplex> CFourierTransformFactoryFFTWFft::CreateRealToComplexPackedFftw(const Options& opts)
{
	// not implemented
	return std::make_shared<CRealToComplex_PackedFftwComplexFFTWFft>(opts.width, opts.height, opts.normalizationFactor);
}

/*virtual*/std::shared_ptr<IRealToComplex_PackedIppComplex> CFourierTransformFactoryFFTWFft::CreateRealToComplexPackedIpp(const Options& opts)
{
	return std::shared_ptr<IRealToComplex_PackedIppComplex>();
}

/*virtual*/std::shared_ptr<IRealToComplex_Complex> CFourierTransformFactoryFFTWFft::CreateRealToComplex(const Options& opts)
{
	return std::make_shared<CRealToComplex_ComplexFFTWFft>(opts.width, opts.height, opts.normalizationFactor);
}

/*virtual*/std::shared_ptr<IRealToComplex_Complex> CFourierTransformFactoryFFTWFft::CreateRealToComplex_Shifted(const Options& opts)
{
	/*if (opts.width % 2 != 0 || opts.height % 2 != 0)
	{
	throw std::runtime_error("Only implemented for even width and height.");
	}*/

	//cufftHandle handle = cuFft_CreatePlan_2dRealToComplex_ZeroFreqInMiddle(opts.width, opts.height);
	return std::shared_ptr<IRealToComplex_Complex>();

}

/*virtual*/std::shared_ptr<IComplexToComplexFwd_Complex> CFourierTransformFactoryFFTWFft::CreateComplexToComplexFwd(const Options& opts)
{
	return std::make_shared<CComplexToComplexFwd_ComplexFFTWFft>(opts.width, opts.height, opts.normalizationFactor);
}

/*virtual*/std::shared_ptr<IComplexToComplexInv_Complex> CFourierTransformFactoryFFTWFft::CreateComplexToComplexInv(const Options& opts)
{
	return std::make_shared<CComplexToComplexInv_ComplexFFTWFft>(opts.width, opts.height, opts.normalizationFactor);
}

/*virtual*/std::shared_ptr<IComplexHermitianToReal> CFourierTransformFactoryFFTWFft::CreateComplexHermitianToReal(const Options& opts)
{
	return std::make_shared<CComplexHermetianToReal_ComplexFFTWFft>(opts.width, opts.height, opts.normalizationFactor);
}